import warnings
warnings.filterwarnings('ignore')

import util_for_output_zh

import pandas as pd
import numpy as np
from pyswarm import pso

import warnings
warnings.filterwarnings('ignore')

import pandas as pd
import numpy as np

from util_set_zh_matplot import plt

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)


# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)

'''额外的约束条件 需要 calc_pass_probability(...) > 0.8
from scipy.stats import norm

def predict_y(week, bmi):
    """根据拟合公式计算模型预测的浓度值Ŷᵢ"""
    return (
        -1.260742 
        + 0.076338 * week + 0.076884 * bmi 
        - 0.005221 * week**2 - 0.002285 * bmi**2 + 0.000798 * week * bmi 
        + 0.000138 * week**3 + 0.000016 * bmi**3 - 0.000063 * week**2 * bmi + 0.000022 * week * bmi**2
    )

def predict_y(week, bmi):
    """根据拟合公式计算模型预测的浓度值Ŷᵢ"""
    return (
        0.0524
        + 0.1736 * week
        + 0.1089 * bmi
        + -0.4638 * week**2
        + 0.0131 * week * bmi
        + -0.2792 * bmi**2
        + 0.3613 * week**3
        + -0.1273 * week**2 * bmi
        + 0.1259 * week * bmi**2
        + 0.1202 * bmi**3
    )
    
def calc_pass_probability(week, bmi, sigma):
    """计算达标概率 P(Yᵢ ≥ 0.04)"""
    y_hat = predict_y(week, bmi)
    z_score = (0.04 - y_hat) / sigma
    pass_prob = 1 - norm.cdf(z_score)
    return pass_prob
'''
from scipy.stats import norm

# def predict_y(week, bmi):
#     """根据拟合公式计算模型预测的浓度值Ŷᵢ"""
#     return (
#         -1.260742 
#         + 0.076338 * week + 0.076884 * bmi 
#         - 0.005221 * week**2 - 0.002285 * bmi**2 + 0.000798 * week * bmi 
#         + 0.000138 * week**3 + 0.000016 * bmi**3 - 0.000063 * week**2 * bmi + 0.000022 * week * bmi**2
#     )
def predict_y(week, bmi):
    """根据拟合公式计算模型预测的浓度值Ŷᵢ"""
    return (
        0.0524
        + 0.1736 * week
        + 0.1089 * bmi
        + -0.4638 * week**2
        + 0.0131 * week * bmi
        + -0.2792 * bmi**2
        + 0.3613 * week**3
        + -0.1273 * week**2 * bmi
        + 0.1259 * week * bmi**2
        + 0.1202 * bmi**3
    )

def calc_pass_probability(week, bmi, sigma):
    """计算达标概率 P(Yᵢ ≥ 0.04)"""
    y_hat = predict_y(week, bmi)
    z_score = (0.04 - y_hat) / sigma
    pass_prob = 1 - norm.cdf(z_score)
    return pass_prob

# 1. 数据加载与预处理
def convert_gestational_week(week_str):
    """快速转换孕周为数值"""
    if 'w+' in week_str:
        w, d = week_str.split('w+')
        return float(w) + float(d)/7
    return float(week_str.split('w')[0])

# 加载数据
df = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')
df = df[['孕妇BMI', '检测孕周', 'Y染色体浓度']]
df = df[df['Y染色体浓度'] > 0.04].copy()
df['检测孕周数值'] = df['检测孕周'].apply(convert_gestational_week)
df = df.drop('检测孕周', axis=1)


# 2. 动态BMI分段
def dynamic_bmi_segmentation(data, n_segments=4, min_ratio=0.1):
    bmi = data['孕妇BMI'].values
    total = len(data)
    min_samples = total * min_ratio
    
    quantiles = np.linspace(0, 1, n_segments + 1)[1:-1]
    cutoffs = np.percentile(bmi, quantiles * 100)
    
    segments = np.digitize(bmi, cutoffs)
    counts = np.bincount(segments, minlength=n_segments)
    
    if not all(counts >= min_samples):
        sorted_bmi = np.sort(bmi)
        cutoffs = []
        for i in range(1, n_segments):
            cutoff_idx = int(min_samples * i)
            cutoffs.append(sorted_bmi[cutoff_idx])
        cutoffs = np.unique(cutoffs)
    
    return np.sort(cutoffs)

n_segments = 4
bmi_cutoffs = dynamic_bmi_segmentation(df, n_segments)
print(f"BMI分界点: {bmi_cutoffs}")

# 预计算每个BMI分段的数据
seg_data_list = []
bmi_ranges = []
for i in range(n_segments):
    if i == 0:
        seg_data = df[df['孕妇BMI'] < bmi_cutoffs[0]]
        bmi_range = (df['孕妇BMI'].min(), bmi_cutoffs[0])
    elif i == n_segments - 1:
        seg_data = df[df['孕妇BMI'] >= bmi_cutoffs[-1]]
        bmi_range = (bmi_cutoffs[-1], df['孕妇BMI'].max())
    else:
        seg_data = df[(df['孕妇BMI'] >= bmi_cutoffs[i-1]) & 
                     (df['孕妇BMI'] < bmi_cutoffs[i])]
        bmi_range = (bmi_cutoffs[i-1], bmi_cutoffs[i])
    seg_data_list.append(seg_data)
    bmi_ranges.append(bmi_range)

# 在数据预处理阶段，为每个BMI分段计算标准差
seg_std_list = []
for seg_data in seg_data_list:
    if len(seg_data) > 1:
        seg_std_list.append(seg_data['Y染色体浓度'].std())
    else:
        seg_std_list.append(0.05)  # 默认值，当分段样本不足时使用


# 3. 定义时间权重函数
def get_time_weight(gestational_week):
    """根据孕周返回权重分数"""
    if gestational_week < 11: return 2.0
    elif 11 <= gestational_week <= 12: return 1.75
    elif 12 < gestational_week <= 15: return 1.5
    elif 15 < gestational_week <= 17: return 1.4
    elif 17 < gestational_week <= 19: return 1.3
    elif 19 < gestational_week <= 21: return 1.2
    elif 21 < gestational_week <= 23: return 1.1
    elif 23 < gestational_week <= 25: return 1.0
    elif 25 < gestational_week <= 27: return 0.5
    else: return 0.0

# 4. 优化的目标函数
def constrained_objective(x):
    """基于时间权重的新目标函数"""
    # 约束检查
    if not (x[0] <= x[1] <= x[2] <= x[3]):
        return -1e6  # 不可行解
    
    if any(x < 10) or any(x > 30):
        return -1e6
    
    # 检查达标概率约束
    for i in range(n_segments):
        bmi_range = bmi_ranges[i]
        avg_bmi = np.mean(bmi_range)  # 使用分段BMI的平均值计算
        # 使用该分段的标准差
        if calc_pass_probability(x[i], avg_bmi, seg_std_list[i]) <= 0.8:
            return -1e6  # 不满足达标概率约束
    
    total_score = 0.0
    diversity_penalty = 0.0
    
    # 计算每个分段的得分
    for i in range(n_segments):
        seg_data = seg_data_list[i]
        if len(seg_data) == 0:
            continue
        
        # 计算该分段在建议时间前的样本
        early_samples = seg_data[seg_data['检测孕周数值'] <= x[i]]
        
        # 计算这些样本的权重得分总和
        segment_score = early_samples['检测孕周数值'].apply(get_time_weight).sum()
        
        # 标准化得分（除以总样本量）
        normalized_score = segment_score / len(seg_data)
        
        total_score += normalized_score
    
    # 多样性惩罚（防止分段过于接近）
    for i in range(1, n_segments):
        diff = x[i] - x[i-1]
        if diff < 1.0:  # 最小允许差异1周
            diversity_penalty += (1.0 - diff) * 100
    
    # 我们希望最大化总得分，所以返回负值（因为PSO默认最小化）
    return -(total_score - diversity_penalty)

# 5. 运行PSO优化
lb = [10, 12, 14, 16]  # 下限
ub = [20, 22, 24, 25]  # 上限

best_x, best_f = pso(
    constrained_objective,
    lb, ub,
    swarmsize=30,
    maxiter=150,
    debug=True,
    phip=0.7,
    phig=0.8,
    minstep=0.1,
    minfunc=0.1
)

# 6. 结果分析
print("\n=== 优化结果 ===")
for i in range(n_segments):
    lower = bmi_cutoffs[i-1] if i > 0 else -np.inf
    upper = bmi_cutoffs[i] if i < n_segments - 1 else np.inf
    seg_data = seg_data_list[i]
    
    print(f"BMI分段 {lower:.1f}-{upper:.1f}:")
    print(f"  最佳检测时间: {best_x[i]:.2f} 周")
    
    # 计算该分段的得分
    early_samples = seg_data[seg_data['检测孕周数值'] <= best_x[i]]
    score = early_samples['检测孕周数值'].apply(get_time_weight).sum()
    normalized_score = score / len(seg_data)
    
    print(f"  得分: {normalized_score:.4f}")
    print(f"  样本量: {len(seg_data)}")
    print(f"  早期样本量: {len(early_samples)}")
    print(f"  早期比例: {len(early_samples)/len(seg_data):.2%}")

# 7. 可视化
plt.figure(figsize=(14, 7))
colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']

for i in range(n_segments):
    seg_data = seg_data_list[i]
    plt.scatter(
        seg_data['检测孕周数值'], 
        seg_data['孕妇BMI'],
        color=colors[i],
        alpha=0.6,
        label=f'BMI分段{i+1}'
    )
    
    plt.axvline(
        x=best_x[i],
        color=colors[i],
        linestyle='--',
        linewidth=2,
        label=f'分段{i+1}最佳时间({best_x[i]:.1f}周)'
    )

plt.xlabel('检测孕周', fontsize=12)
plt.ylabel('孕妇BMI', fontsize=12)
plt.title('基于时间权重评分的最佳检测时间点优化', fontsize=14)
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('ques2_optimized_detection_time_with_score.png', dpi=300, bbox_inches='tight')
plt.close()

# 8. 保存结果
results = []
for i in range(n_segments):
    lower = bmi_cutoffs[i-1] if i > 0 else df['孕妇BMI'].min()
    upper = bmi_cutoffs[i] if i < n_segments - 1 else df['孕妇BMI'].max()
    seg_data = seg_data_list[i]
    early_samples = seg_data[seg_data['检测孕周数值'] <= best_x[i]]
    score = early_samples['检测孕周数值'].apply(get_time_weight).sum()
    normalized_score = score / len(seg_data)
    
    results.append({
        'BMI范围': f"{lower:.1f}-{upper:.1f}",
        '最佳检测时间(周)': best_x[i],
        '得分': normalized_score,
        '样本量': len(seg_data),
        '早期样本量': len(early_samples),
        '早期比例': len(early_samples)/len(seg_data)
    })

results_df = pd.DataFrame(results)
results_df.to_csv('ques2_optimized_detection_results.csv', index=False)
print("\n优化结果已保存到 ques2_optimized_detection_results.csv")